#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import tempfile
import re
import logging
import uuid
from typing import Dict, List, Any, BinaryIO, Tuple, Optional
from pathlib import Path
import shutil
import io
from datetime import datetime

from marker.converters.pdf import PdfConverter
from marker.models import create_model_dict
from marker.output import text_from_rendered
from marker.config.parser import ConfigParser

from document_processor.engines.base_engine import BaseEngine

logger = logging.getLogger(__name__)

class MarkerEngine(BaseEngine):
    """使用 Marker 处理文档的引擎
    
    这是一个高级文档处理引擎，基于 VikParuchuri 的 Marker 项目。
    支持 PDF、Word、PPT 等多种格式的文档转换，具有更好的布局分析和内容提取能力。
    """
    
    def __init__(self, config: Dict[str, Any]):
        """初始化 Marker 处理引擎
        
        Args:
            config: 配置字典，包含 Marker 相关的配置
        """
        super().__init__(config)
        
        # 初始化 Marker 配置
        self.marker_config = config.get("marker", {})
        
        # 默认配置
        default_config = {
            "output_format": "markdown",
            "use_llm": self.marker_config.get("use_llm", False),
            "force_ocr": self.marker_config.get("force_ocr", False),
            "strip_existing_ocr": self.marker_config.get("strip_existing_ocr", False),
            "debug": False,  # 开启调试模式会输出更多信息
        }
        
        # 合并用户配置
        self.merged_config = {**default_config, **self.marker_config}
        logger.info(f"Marker 配置: {self.merged_config}")
        
        # 创建配置解析器进行验证
        try:
            config_parser = ConfigParser(self.merged_config)
            # 验证配置是否有效
            _ = config_parser.generate_config_dict()
            _ = config_parser.get_processors()
            _ = config_parser.get_renderer()
            _ = config_parser.get_llm_service()
            logger.info("已完成 Marker 配置验证")
        except Exception as e:
            logger.error(f"初始化 Marker 配置失败: {str(e)}")
            raise

    async def convert_to_markdown(
        self, 
        file_content: BinaryIO, 
        file_name: str,
        image_option: str = "obs",
        work_dir: Optional[str] = None,
        md_filename: Optional[str] = None
    ) -> Dict[str, Any]:
        """使用 Marker 将文档转换为 Markdown 格式
        
        Args:
            file_content: 文件内容（二进制流）
            file_name: 文件名
            image_option: 图片处理选项，可选值："obs"（上传到OBS）或"remove"（移除图片）
            work_dir: 工作目录路径，如果提供则使用此目录，否则创建临时目录
            md_filename: 生成的Markdown文件名（不含扩展名），用于命名图片目录
            
        Returns:
            包含以下键的字典:
                - text: 提取的纯文本
                - markdown: 转换后的Markdown文本
                - images: 图片信息列表
                - temp_dir: 临时目录路径
        """
        # 验证图片处理选项
        if image_option not in ["obs", "remove"]:
            logger.warning(f"不支持的图片处理选项: {image_option}，使用 obs 选项")
            image_option = "obs"
        
        # 检查文件格式
        file_ext = os.path.splitext(file_name)[1].lower()
        supported_formats = ['.pdf', '.docx', '.pptx', '.xlsx', '.html', '.epub']
        if file_ext not in supported_formats:
            raise ValueError(f"Marker引擎不支持该格式: {file_ext}")
            
        # 使用提供的工作目录或创建临时目录
        temp_dir = work_dir if work_dir else tempfile.mkdtemp()
        if not os.path.exists(temp_dir):
            os.makedirs(temp_dir)
            
        # 确保output目录存在
        output_dir = os.path.join(temp_dir, "output")
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            
        temp_file_path = os.path.join(temp_dir, file_name)
        
        try:
            # 保存上传的文件到临时位置
            file_content_bytes = file_content.read()
            with open(temp_file_path, "wb") as f:
                f.write(file_content_bytes)
            
            # 创建图片目录
            if md_filename:
                images_dir_name = f"{md_filename}_images"
            else:
                dir_id = uuid.uuid4().hex[:8]
                images_dir_name = f"doc_images_{dir_id}"
                
            images_dir = os.path.join(output_dir, images_dir_name)
            os.makedirs(images_dir, exist_ok=True)
            
            logger.info(f"开始使用 Marker 转换文档: {file_name}")
            
            # 设置是否提取图片
            if image_option == "remove":
                # 创建新的配置，包含禁用图片提取
                new_config = {**self.merged_config, "disable_image_extraction": True}
            else:
                # 确保图片提取被启用
                new_config = {**self.merged_config, "disable_image_extraction": False}
            
            # 重新创建配置解析器
            config_parser = ConfigParser(new_config)
            
            # 使用正确的方式创建转换器实例
            converter = PdfConverter(
                config=config_parser.generate_config_dict(),
                artifact_dict=create_model_dict(),
                processor_list=config_parser.get_processors(),
                renderer=config_parser.get_renderer(),
                llm_service=config_parser.get_llm_service()
            )
            
            try:
                # 按照官方文档方式直接调用转换器处理文档
                rendered = converter(temp_file_path)
                
                # 从渲染结果中提取内容 - 注意这里只有两个返回值，text 即为 markdown 内容
                text,md, images_data = text_from_rendered(rendered)

                markdown = text
                logger.info(f"成功从渲染结果中提取内容，Markdown字符数: {len(text)}")
                
            except Exception as e:
                logger.error(f"Marker处理过程出错: {str(e)}")
                import traceback
                logger.error(f"详细错误: {traceback.format_exc()}")
                
                text = f"转换失败: {str(e)}"
                markdown = f"# {Path(file_name).stem}\n\n无法使用Marker转换文档。\n\n错误: {str(e)}"
                images_data = {}
            
            # 检查markdown内容是否有效
            if not markdown or len(markdown.strip()) < 5 or markdown.strip() == "md":
                logger.error(f"Marker返回的Markdown内容无效: '{markdown}'")
                markdown = f"# {Path(file_name).stem}\n\n无法使用Marker转换获取有效内容。"
                text = f"无法获取有效内容。"
                images_data = {}
            
            # 处理图片 - 注意images_data现在是字典格式，键是图片文件名，值是PIL.Image.Image对象
            processed_images = []
            if image_option != "remove" and images_data:
                for img_filename, img_obj in images_data.items():
                    try:
                        # 确保是PIL图像对象
                        if not hasattr(img_obj, 'save'):
                            logger.warning(f"跳过非PIL图像对象: {type(img_obj)}")
                            continue
                        
                        # 获取图片格式
                        img_format = os.path.splitext(img_filename)[1].lower().replace('.', '')
                        if not img_format or img_format not in ['jpg', 'jpeg', 'png', 'gif']:
                            img_format = 'jpeg'  # 默认使用JPEG格式
                        
                        # 生成图片保存路径
                        img_path = os.path.join(images_dir, img_filename)
                        
                        # 保存图片
                        img_buffer = io.BytesIO()
                        img_obj.save(img_buffer, format=img_format.upper())
                        img_data = img_buffer.getvalue()
                        
                        with open(img_path, "wb") as img_file:
                            img_file.write(img_data)
                        
                        # 生成图片信息
                        image_info = {
                            "filename": img_filename,
                            "format": img_format,
                            "data": img_data,
                            "dir_name": images_dir_name,
                            "path": img_path
                        }
                        
                        processed_images.append(image_info)
                        logger.info(f"成功处理并保存图片: {image_info['filename']}")
                        
                    except Exception as e:
                        logger.error(f"处理图片失败: {str(e)}")
                        import traceback
                        logger.error(f"处理图片详细错误: {traceback.format_exc()}")
                        continue
            
            # 保存 Markdown 到临时文件
            output_file = os.path.join(output_dir, f"{Path(file_name).stem}.md")
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(markdown)
            
            # 日志输出
            logger.info(f"生成的Markdown长度: {len(markdown)} 字符")
            logger.info(f"处理了 {len(processed_images)} 张图片")
            
            return {
                "text": text,
                "markdown": markdown,
                "images": processed_images,
                "temp_dir": temp_dir,
                "output_md": output_file,
                "images_dir": images_dir,
                "md_filename": Path(file_name).stem
            }
            
        except Exception as e:
            logger.error(f"处理文件出错: {str(e)}")
            import traceback
            logger.error(f"详细错误: {traceback.format_exc()}")
            
            # 创建基本的错误MD文件
            output_file = os.path.join(output_dir, f"{Path(file_name).stem}.md")
            error_md = f"# 转换错误\n\n无法处理文件 {file_name}\n\n错误: {str(e)}"
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(error_md)
                
            return {
                "text": f"转换错误: {str(e)}",
                "markdown": error_md,
                "images": [],
                "temp_dir": temp_dir,
                "output_md": output_file,
                "images_dir": os.path.join(output_dir, "images"),
                "md_filename": Path(file_name).stem
            }

    async def extract_images(
        self,
        file_content: BinaryIO,
        file_name: str,
        work_dir: Optional[str] = None
    ) -> Dict[str, Any]:
        """从文档中提取图片
        
        Args:
            file_content: 文件内容（二进制流）
            file_name: 文件名
            work_dir: 工作目录路径，如果提供则使用此目录，否则创建临时目录
            
        Returns:
            包含以下键的字典:
                - images: 图片信息列表
                - temp_dir: 临时目录路径
        """
        # 使用提供的工作目录或创建临时目录
        temp_dir = work_dir if work_dir else tempfile.mkdtemp()
        if not os.path.exists(temp_dir):
            os.makedirs(temp_dir)
            
        # 确保output目录存在
        output_dir = os.path.join(temp_dir, "output")
        images_dir = os.path.join(output_dir, "images")
        os.makedirs(images_dir, exist_ok=True)
            
        temp_file_path = os.path.join(temp_dir, file_name)
        
        try:
            # 保存上传的文件到临时位置
            file_content_bytes = file_content.read()
            with open(temp_file_path, "wb") as f:
                f.write(file_content_bytes)
            
            logger.info(f"开始从文档提取图片: {file_name}")
            
            # 确保图片提取被启用
            new_config = {**self.merged_config, "disable_image_extraction": False}
            
            # 重新创建配置解析器
            config_parser = ConfigParser(new_config)
            
            # 使用正确的方式创建转换器实例
            converter = PdfConverter(
                config=config_parser.generate_config_dict(),
                artifact_dict=create_model_dict(),
                processor_list=config_parser.get_processors(),
                renderer=config_parser.get_renderer(),
                llm_service=config_parser.get_llm_service()
            )
            
            try:
                # 按照官方文档方式直接调用转换器处理文档
                rendered = converter(temp_file_path)
                
                # 从渲染结果中提取内容 - 只关注图片
                _, images = text_from_rendered(rendered)
                
            except Exception as e:
                logger.error(f"提取图片过程出错: {str(e)}")
                import traceback
                logger.error(f"详细错误: {traceback.format_exc()}")
                images = {}
            
            # 处理图片 - images是一个字典，键是图片文件名，值是PIL.Image.Image对象
            processed_images = []
            if images:
                for img_filename, img_obj in images.items():
                    try:
                        # 确保是PIL图像对象
                        if not hasattr(img_obj, 'save'):
                            logger.warning(f"跳过非PIL图像对象: {type(img_obj)}")
                            continue
                        
                        # 获取图片格式
                        img_format = os.path.splitext(img_filename)[1].lower().replace('.', '')
                        if not img_format or img_format not in ['jpg', 'jpeg', 'png', 'gif']:
                            img_format = 'jpeg'  # 默认使用JPEG格式
                        
                        # 生成图片保存路径
                        img_path = os.path.join(images_dir, img_filename)
                        
                        # 保存图片
                        img_buffer = io.BytesIO()
                        img_obj.save(img_buffer, format=img_format.upper())
                        img_data = img_buffer.getvalue()
                        
                        with open(img_path, "wb") as f:
                            f.write(img_data)
                        
                        # 生成图片信息
                        image_info = {
                            "filename": img_filename,
                            "format": img_format,
                            "data": img_data,
                            "path": img_path
                        }
                        
                        processed_images.append(image_info)
                        logger.info(f"成功提取并保存图片: {image_info['filename']}")
                        
                    except Exception as e:
                        logger.error(f"处理图片失败: {str(e)}")
                        continue
            
            logger.info(f"成功从文档中提取 {len(processed_images)} 张图片")
            return {
                "images": processed_images,
                "temp_dir": temp_dir
            }
            
        except Exception as e:
            logger.error(f"提取图片时出错: {str(e)}")
            import traceback
            logger.error(f"详细错误: {traceback.format_exc()}")
            return {
                "images": [],
                "temp_dir": temp_dir
            }

    async def handle_images(
        self, 
        images: List[Dict[str, Any]], 
        markdown_text: str,
        uploader,
        md_filename: Optional[str] = None
    ) -> str:
        """处理图片（上传到OBS并更新Markdown中的图片引用）
        
        Args:
            images: 图片信息列表
            markdown_text: Markdown文本
            uploader: OBS上传器实例
            md_filename: Markdown文件名（不含扩展名），用于生成OBS路径
            
        Returns:
            更新后的Markdown文本（图片引用已替换为OBS URL）
        """
        if not images:
            logger.info("没有图片需要处理")
            return markdown_text
        
        # 检查markdown_text是否有效
        if not markdown_text or len(markdown_text.strip()) < 5:
            logger.warning(f"待处理的Markdown内容无效: '{markdown_text}'")
            return markdown_text
        
        # 如果没有提供OBS上传器，则无法处理图片
        if not uploader or not hasattr(uploader, 'upload'):
            logger.warning("未提供有效的OBS上传器，无法处理图片")
            return markdown_text
        
        logger.info(f"开始处理 {len(images)} 张图片")
        
        # 检查是否有图片引用模式: ![](filename.jpg)
        image_pattern = r'!\[\]\((.*?)\)'
        referenced_images = re.findall(image_pattern, markdown_text)
        
        if not referenced_images:
            logger.warning("Markdown中没有找到图片引用")
            
        try:
            # 图片URL映射表
            image_urls = {}
            
            # 生成OBS路径前缀，按照日期构建
            date_str = datetime.now().strftime("%Y%m%d")
            obs_path_prefix = f"documents/{date_str}"
            
            # 为每个图片创建一个文件名到对象的映射
            image_map = {image['filename']: image for image in images}
            
            # 首先处理已在Markdown中引用的图片
            for img_ref in referenced_images:
                # 检查图片是否在我们的图片列表中
                if img_ref in image_map:
                    try:
                        image = image_map[img_ref]
                        
                        # 准备上传数据
                        upload_data = {
                            "filename": image['filename'],
                            "data": image['data'],
                            "content_type": f"image/{image['format']}"
                        }
                        
                        # 构建OBS路径
                        safe_doc_name = re.sub(r'[^\w\-\.]', '_', md_filename if md_filename else "doc")
                        # 保留原始图片名称
                        obs_path = f"{obs_path_prefix}/{safe_doc_name}_{image['filename']}"
                        upload_data["object_key"] = obs_path
                        
                        # 上传图片
                        img_url = await uploader.upload(upload_data)
                        if not img_url:
                            logger.warning(f"图片 {image['filename']} 上传返回空URL")
                            continue
                            
                        # 记录URL
                        image_urls[img_ref] = img_url
                        logger.info(f"图片上传成功: {img_url}")
                        
                        # 替换Markdown中的图片标记
                        markdown_text = markdown_text.replace(f"![]({img_ref})", f"![图片]({img_url})")
                    
                    except Exception as e:
                        logger.error(f"处理图片 {img_ref} 失败: {str(e)}")
                        continue
                else:
                    logger.warning(f"引用的图片 {img_ref} 不在提取的图片列表中")
            
            # 处理未在Markdown中引用但在图片列表中的图片
            unreferenced_images = [img for img in images if img['filename'] not in referenced_images]
            if unreferenced_images:
                logger.info(f"处理 {len(unreferenced_images)} 张未引用的图片")
                
                for i, image in enumerate(unreferenced_images):
                    try:
                        # 准备上传数据
                        upload_data = {
                            "filename": image['filename'],
                            "data": image['data'],
                            "content_type": f"image/{image['format']}"
                        }
                        
                        # 构建OBS路径
                        safe_doc_name = re.sub(r'[^\w\-\.]', '_', md_filename if md_filename else "doc")
                        obs_path = f"{obs_path_prefix}/{safe_doc_name}_{image['filename']}"
                        upload_data["object_key"] = obs_path
                        
                        # 上传图片
                        img_url = await uploader.upload(upload_data)
                        if not img_url:
                            logger.warning(f"图片 {image['filename']} 上传返回空URL")
                            continue
                            
                        # 记录URL
                        image_urls[image['filename']] = img_url
                        logger.info(f"未引用图片上传成功: {img_url}")
                        
                        # 如果图片没有在Markdown中引用，可以选择在文档末尾添加引用
                        # 这里可以根据需要注释或启用
                        # markdown_text += f"\n\n![附加图片 {i+1}]({img_url})\n"
                        
                    except Exception as e:
                        logger.error(f"处理未引用图片 {image['filename']} 失败: {str(e)}")
                        continue
            
            return markdown_text
                
        except Exception as e:
            logger.error(f"处理图片出错: {str(e)}")
            import traceback
            logger.error(f"详细错误: {traceback.format_exc()}")
            return markdown_text 